from numpy.ma.extras import row_stack

import rclpy
from rclpy.serialization import deserialize_message
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions
from sensor_msgs.msg import PointCloud2, Imu

def read_bag_file(bag_path: str):
    # 初始化ROS 2 Python客户端
    rclpy.init()

    # 配置存储选项
    storage_options = StorageOptions(uri=bag_path, storage_id='sqlite3')
    converter_options = ConverterOptions('', '')  # 输入输出格式通常为空字符串

    # 创建读取器实例
    reader = SequentialReader()
    reader.open(storage_options, converter_options)

    # 获取话题类型映射
    topic_types = reader.get_all_topics_and_types()
    type_map = {topic.name: topic.type for topic in topic_types}
    # {'/unilidar/imu': 'sensor_msgs/msg/Imu', '/unilidar/cloud': 'sensor_msgs/msg/PointCloud2'}

    # 遍历消息
    row_count = 0
    while reader.has_next():
        row_count += 1
        topic, data, timestamp = reader.read_next()
        if 'PointCloud2' in type_map[topic]:  # 示例：筛选点云数据
            msg = deserialize_message(data, PointCloud2)
            print(f"{row_count:6d})Timestamp: {timestamp}, Type: {topic}, Points: {msg.width * msg.height}")
        else:

            msg = deserialize_message(data, Imu)
            # sensor_msgs.msg.Imu(header=std_msgs.msg.Header(stamp=builtin_interfaces.msg.Time(sec=1695357725, nanosec=165899515), frame_id='unilidar_imu'), orientation=geometry_msgs.msg.Quaternion(x=0.03746518865227699, y=0.6997672915458679, z=-0.039229828864336014, w=0.7099371552467346), orientation_covariance=array([0., 0., 0., 0., 0., 0., 0., 0., 0.]), angular_velocity=geometry_msgs.msg.Vector3(x=0.002303345361724496, y=0.006947698071599007, z=0.0028555006720125675), angular_velocity_covariance=array([0., 0., 0., 0., 0., 0., 0., 0., 0.]), linear_acceleration=geometry_msgs.msg.Vector3(x=-9.709754943847656, y=-0.10081425309181213, z=0.3058036267757416), linear_acceleration_covariance=array([0., 0., 0., 0., 0., 0., 0., 0., 0.]))  # noqa
            print(f"{row_count:6d})Timestamp: {timestamp}, Type: {topic}, Imu: {msg.orientation}, {msg.angular_velocity}")

    # 清理资源
    reader.close()
    rclpy.shutdown()


if __name__ == '__main__':
    bag_file = "/gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3"  # 替换为实际路径
    read_bag_file(bag_file)
